package me.winter.imgcmps;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.TreeMap;

import javax.swing.plaf.basic.BasicInternalFrameTitlePane.MaximizeAction;

import org.omg.CORBA.PRIVATE_MEMBER;





/**
 * Main/launch program for Huff assignment. A better
 * comment than this is warranted.
 * @author Winter Lu
 *
 */
public class Huff {
	
	private int[] in;
	private int width, height;
	private int offset;
	
	private final int INTERNAL_NODE_VALUE = 512;
	private final int PSEUDO_EOF = 0x1119;
	private boolean DEBUG = false;
	private boolean PRINT_HEADER_SIZE = true;
	
	public Huff(int[][] input) {
		
		// Convert two dimension data to one dimension
		// and get the minimum value
		int mininum = 256;
		width = input[0].length;
		height = input.length;
		
		in = new int[height * width];
		for (int i = 0; i < height; ++i) {
			for (int j = 0; j < width; ++j) {
				in[i * width + j] = input[i][j];
				if (input[i][j] < mininum)
					mininum = input[i][j];
			}
		}
		debugPrintln("The minimum value is " + mininum);
		
		// set the offset
		offset = mininum < 0 ? -mininum : 0;
		debugPrintln("The offset = " + offset);
	}
	
	public Huff(int[] input) {
		
		in = new int[input.length];
		
		System.arraycopy(input, 0, in, 0, input.length);
		
		int minimum = 256;
		for (int i = 0; i < input.length; ++i) {
			if (input[i] < minimum)
				minimum = input[i];
		}
		
		offset = minimum < 0 ? -minimum : 0;
		debugPrintln("The offset = " + offset);
	}
	
	private void getFrequencyTable(Map<Integer, Integer> frequencies) {
		
		int maxFreq = 0, mostFreqVal = 256;
		for (int v : in) {
			if (frequencies.containsKey(v + offset)) {
				frequencies.put(v + offset, frequencies.get(v + offset) + 1);
			}
			else {
				frequencies.put(v + offset, 1);
			}
			
			if (frequencies.get(v + offset) > maxFreq) {
				maxFreq = frequencies.get(v + offset);
				mostFreqVal = v + offset;
			}
		}
		frequencies.put(PSEUDO_EOF, 1);
		
		debugPrintln("The size of frequency table is " + frequencies.size());
	}

	/**
	 * Inner private class
	 * @author Winter
	 *
	 */
	private class TreeNode implements Comparable<TreeNode> {
		public int symbol;
		public int weight;
		public TreeNode zero;
		public TreeNode one;
		
		TreeNode(int val, int wght) {
			symbol = val;
			weight = wght;
			zero = null;
			one = null;
		}

		@Override
		public int compareTo(TreeNode rhs) {
			return this.weight - rhs.weight;
		}
	}
	
	private TreeNode buildEncodingTree(Map<Integer, Integer> frequencies) {
		PriorityQueue<TreeNode> pq = new PriorityQueue<TreeNode>();	
		
		for (Integer key : frequencies.keySet()) {
			pq.add(new TreeNode(key, frequencies.get(key)));
		}
		
		while (pq.size() > 1) {
			TreeNode min = pq.poll();
			TreeNode nextToMin = pq.poll();
			
			TreeNode parent = new TreeNode(INTERNAL_NODE_VALUE,
					min.weight + nextToMin.weight);
			parent.zero = min;
			parent.one = nextToMin;
			pq.add(parent);
		}
		
		return pq.poll();
	}
	
	
	private void genEncodingTable(TreeNode root, String codeword, Map<Integer, String> codes) {
    	if (root == null)
    		return;
    	// Base case: root is a leaf
    	if (root.symbol != INTERNAL_NODE_VALUE) {
    		codes.put(root.symbol, codeword);
    	}
    	// Recursive case
    	else {
    		genEncodingTable(root.zero, codeword + "0", codes);
    		genEncodingTable(root.one, codeword + "1", codes);
    	}
	}
	
	
	private void writeImageHeader(Map<Integer, Integer> frequencies,
			BitOutputStream bitOs) throws IOException {
		
		if (! frequencies.containsKey(PSEUDO_EOF)) {
			throw new IOException("File Format Error");
		}
		
		bitOs.writeBits(15, frequencies.size() - 1);
		
		for (Integer key : frequencies.keySet()) {
			if (key != PSEUDO_EOF) {
				bitOs.writeBits(9, key);
				bitOs.writeBits(32, frequencies.get(key));
			}
		}

		if (PRINT_HEADER_SIZE ) {
			DEBUG = true;
			int totalBits = 15 + (frequencies.size() - 1) * (9 + 32);
			debugPrintln("Header size: " + totalBits + "bits" +
					"(" + totalBits / 8 + "bytes)");
			DEBUG = false;
		}
	}
	
	private void encodeImage(int[] in, Map<Integer, String> codes,
			BitOutputStream bitOs) {
		debugPrintln("Encoding image...");
		for (int i = 0; i < in.length; ++i) {
			String codeword = codes.get(in[i] + offset);
			debugPrintln("Writing codeword " + codeword);
			for (int j = 0; j < codeword.length(); ++j)
				bitOs.writeBits(1, codeword.charAt(j) - '0');
		}
		
		String codeword = codes.get(PSEUDO_EOF);
		debugPrintln("Writing codeword " + codeword);
		for (int i = 0; i < codeword.length(); ++i)
			bitOs.writeBits(1, codeword.charAt(i) - '0');
	}
	
	public void compress(String outfilename) {
		Map<Integer, Integer> frequencies = new TreeMap<Integer, Integer>();
		
		getFrequencyTable(frequencies);
		
		Map<Integer, String> codes = new TreeMap<Integer, String>();
		TreeNode treeRoot = buildEncodingTree(frequencies);
		
		genEncodingTable(treeRoot, "", codes);
		debugPrintln("Size of encoding table: " + codes.size());
		debugPrintln("The encoding table is");
		for (Integer key : codes.keySet())
			debugPrintln(key + ": " + codes.get(key));
		
		
		String prefix = "AllImages/Compression/compressed/";
		String compressedFileName = prefix + outfilename;
		File outfile = new File(compressedFileName + ".hf");
		
		BitOutputStream bos = null;
		try {
			bos = new BitOutputStream(
					new BufferedOutputStream(
							new FileOutputStream(outfile)));
			writeImageHeader(frequencies, bos);
			
			debugPrintln("Contents of input");
			for (int i = 0; i < in.length; ++i)
				debugPrint(in[i] + " ");
			debugPrintln("");
			
			encodeImage(in, codes, bos);
			bos.flush();
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			bos.close();
		}
	}
	
	public int[] uncompress(String infilename) {
		
		debugPrintln("Uncompressing...");
		
		String prefix = "AllImages/Compression/compressed/";
		String compressedFileName = prefix + infilename;
		File infile = new File(compressedFileName + ".hf");
		
		BitInputStream bitIs = new BitInputStream(infile);
		Map<Integer, Integer> frequencies = readImageHeader(bitIs);
		TreeNode root = buildEncodingTree(frequencies);
		
		Map<Integer, Integer> freq = new TreeMap<Integer, Integer>();
		getFrequencyTable(freq);
		
		for (Integer key : freq.keySet()) {
			if (frequencies.containsKey(key) &&
					frequencies.get(key).equals(freq.get(key)))
			{
			
			} else {
				debugPrintln("Frequency table doesn't match");
				break;
			}
		}		
		
		
		int[] out = decodeImage(bitIs, root);
		
		
		return out;
	}
	
	public int[][] to2DArray(String infilename) {
		int[] out = uncompress(infilename);
		
		// Convert imgData to two dimension
		int[][] imgData = new int[height][width];
		for (int i = 0; i < height; ++i) {
			for (int j = 0; j < width; ++j) {
				imgData[i][j] = out[i * width + j];
			}
		}
		
		return imgData;
	}
	
	private Map<Integer, Integer> readImageHeader(BitInputStream bitIs) {
		Map<Integer, Integer> freqTable = new TreeMap<Integer, Integer>();
		try {
			int sizeTable = bitIs.readBits(15);
			for (int i = 0; i < sizeTable; ++i) {
				int symbol = bitIs.readBits(9);
				int weight = bitIs.readBits(32);
				freqTable.put(symbol, weight);
			}
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		freqTable.put(PSEUDO_EOF, 1);
		
		return freqTable;
	}
	
	private int[] decodeImage(BitInputStream bis, TreeNode root) {
		int[] out = new int[in.length];
		try {
			int i = 0;
			TreeNode cur = root;
			debugPrintln("Reading contents of uncompressed...");
			int bit;
			while ((bit = bis.readBits(1)) != -1) {
				if (bit == 0) {
					debugPrint("0");
					cur = cur.zero; 
				} else if (bit == 1) {
					debugPrint("1");
					cur = cur.one;
				}
				
				if (cur.symbol != INTERNAL_NODE_VALUE) {
					if (cur.symbol == PSEUDO_EOF)
						break;
					debugPrintln("");
					out[i++] = cur.symbol - offset;
					cur = root;
				}
			}
			// debugPrintln();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		return out;
	}
	
	private void debugPrint(String msg) {
		if (DEBUG ) {
			System.out.print(msg);
		}
	}
	
	private void debugPrintln(String msg) {
		if (DEBUG) {
			System.out.println(msg);
		}
	}
}